/* ImageHandler.cpp
-----------------------------------------
By: 	Wei YE ( ye0003ei@e.ntu.edu.sg )

Date:	10/Oct/2012
-----------------------------------------
*/

#include "ImageHandler.h"

using namespace cv;

//--------------------------------------------------//
//	  Functions For Loading and Generating Images	//
//--------------------------------------------------//

void ImageHandler::LoadTheImage(Mat& image, ImageType imagetype )
{
	Mat img = Mat(image);
	H_ = image.rows;
	W_ = image.cols;

	if(imagetype == GRAYSCALE)
		D_ = 3;
	else
		D_ = 5;

	// allocate memory for data
	data_ = new float[H_*W_*D_]; // nPoints * nDim
	float* tmp = new float[3];

	// load data
	if(D_ == 5)
	{
		for(int i = 0; i < H_; i++)
		{
			for(int j = 0; j < W_; j++)
			{
				data_[(i*W_+j)*D_] = (float)i/SS_;
				data_[(i*W_+j)*D_ + 1] = (float)j/SS_;

				RGBtoLUV(img.ptr<Vec3b>(i,j), tmp);
				data_[(i*W_+j)*D_ + 2] = tmp[0]/SR_;
				data_[(i*W_+j)*D_ + 3] = tmp[1]/SR_;  
				data_[(i*W_+j)*D_ + 4] = tmp[2]/SR_;  
			}
		}
	}
	else if(D_ == 3)
	{
		uv_ = new float[H_*W_*2];
		for(int i = 0; i < H_; i++)
		{
			for(int j = 0; j < W_; j++)
			{
				data_[(i*W_+j)*D_] = (float)i/SS_;
				data_[(i*W_+j)*D_ + 1] = (float)j/SS_;

				RGBtoLUV(img.ptr<Vec3b>(i,j), tmp);
				data_[(i*W_+j)*D_ + 2] = tmp[0]/SR_;  // we only use the L* component
				uv_[(i*W_+j)*2] = tmp[1];             // store the u* and v* componment
				uv_[(i*W_+j)*2 + 1] = tmp[2];
			}
		}
	}
	// define 8-connected neighborhood
	neigh[0]	= 1;
	neigh[1]	= 1 - W_;
	neigh[2]	= -W_;
	neigh[3]	= -(1 + W_);
	neigh[4]	= -1;
	neigh[5]	= W_ -1;
	neigh[6]	= W_;
	neigh[7]	= W_+1;

	// allocate memory for some variables
	modes_    = new float[H_*W_*(D_-2)];
	labels_   = new int[H_*W_];
	idxTable_ = new int[H_*W_];

	delete []tmp;
}

//--------------------------------------------------//
//	      Functions For Mean Shift Algorithm  	    //
//--------------------------------------------------//
void ImageHandler::Segment()
{
	if( !data_ )
		return;

	MS_.LoadPoints(data_, H_*W_, D_);	
	MS_.GridMsFilter(H_, W_, SS_, SR_);

	Merge();
	Eliminate();
}

void ImageHandler::Merge()
{
	float* msout = MS_.getConvPts();
	if(!msout){
		return;
	}

	int i;
	for(i = 0; i < W_*H_; i++)
		labels_[i]	 = -1;

	int label = -1;
	for(i = 0; i < W_*H_; i++)
	{
		//if this region has not yet been labeled - label it
		if(labels_[i] < 0)
		{
			// Assign new label to this region
			labels_[i] = ++label;

			// Add this region to the region list
			rList_.push_back(new Region(label));

			// Copy region color into modes
			if(D_ == 3)
				modes_[label] = (msout[3*i + 2]);
			else if(D_ == 5){
				modes_[label*3 + 0] = (msout[5*i + 2]);
				modes_[label*3 + 1] = (msout[5*i + 3]);
				modes_[label*3 + 2] = (msout[5*i + 4]);
			}
			// label the neighbor pixel
			Fill(i, label, msout);
		}
	}
}

void ImageHandler::Fill(int pt, int label, float* msout)
{
	int	i, k, neighLoc, neighborsFound;
	float dist;

	int	index	 = 0;
	idxTable_[0] = pt;
	rList_.back()->ptInMode_.push_back(pt);

	while(true)
	{
		neighborsFound	= 0;

		//check the eight connected neighbors at regionLoc -
		//if a pixel has similar color to that located at 
		//"pt" then declare it as part of this region
		for(i = 0; i < 8; i++)
		{
			//check bounds and if neighbor has been already labeled
			neighLoc = pt + neigh[i];
			if((neighLoc >= 0)&&(neighLoc < W_*H_)&&(labels_[neighLoc] < 0))
			{
				dist = 0;			
				for(k = 0; k < D_; k++)
					dist += (msout[pt*D_+k]-msout[neighLoc*D_+k])*(msout[pt*D_+k]-msout[neighLoc*D_+k]);
			
				if(dist < 0.5)
				{				
					labels_[neighLoc] = label;      //assign label to neighbor i					
					idxTable_[++index] = neighLoc;	//place index of neighbor i onto the index tabel buffer
					neighborsFound	= 1;
					rList_.back()->ptInMode_.push_back(neighLoc); // add to the pixel list of this region
					for(k = 0; k < D_-2; k++)
						modes_[label*(D_- 2) + k] = (modes_[label*(D_- 2) + k] + msout[neighLoc*D_+k + 2])/2;
				}
			}
		}

		// check the indexTable to see if there are any more entries to be explored 
		// if so explore them, otherwise exit the loop 
		if(neighborsFound)
			pt	= idxTable_[index];
		else if (index > 1)
			pt	= idxTable_[--index];
		else
			break; //fill complete
	}
}

void ImageHandler::Eliminate()
{	
	if( rList_.size() == 0)
		return;

	int i, label, new_label;
	float el, dist, minDist;
	list<Region*>::iterator region_itr, region_itr_r;
	vector<int>::iterator pt_itr;
	int numRegion = rList_.size();
	int* adjVec = new int[numRegion];

	// for every existing region, check their size
	region_itr = rList_.begin();	
	while( region_itr != rList_.end())
	{
		// if size < M, find all adjacent regions of this region
		if((int)(*region_itr)->ptInMode_.size() < M_)
		{
			label = (*region_itr)->ID_;
			memset(adjVec, 0, numRegion*sizeof(int));
			for(pt_itr = (*region_itr)->ptInMode_.begin(); pt_itr != (*region_itr)->ptInMode_.end(); pt_itr++)
			{
				if((*pt_itr)%W_ < W_ && label != labels_[*(pt_itr) + 1])
					adjVec[labels_[*(pt_itr) + 1]]++;
			}

			// find the closest(in color) neighbor region
			for(i = 0; i < numRegion; i++)
			{
				minDist = 1e+5;
				if( adjVec[i] != 0)
				{
					el = modes_[(D_-2)*label] - modes_[(D_-2)*i];
					dist = el*el;
					if(D_ == 5)
					{
						el = modes_[3*label+1] - modes_[3*i+1];
						dist += (el*el);
						el = modes_[3*label+2] - modes_[3*i+2];
						dist += (el*el);
					}
					if( dist < minDist)
					{
						minDist = dist;
						new_label = i;
					}
				}
			}
			// add current region to the closest adjacent region	
			for(region_itr_r = rList_.begin(); region_itr_r != rList_.end(); region_itr_r++)
				if( (*region_itr_r)->ID_ == new_label ) 
					break;

			// re-label pixels in current region	
			for(pt_itr = (*region_itr)->ptInMode_.begin(); pt_itr != (*region_itr)->ptInMode_.end(); pt_itr++)
			{
				(*region_itr_r)->ptInMode_.push_back(*pt_itr);
				labels_[*pt_itr] = new_label;
			}
		}
		region_itr++;
	}
	delete[] adjVec;
}


//--------------------------------------------------//
//			Functions For Color Conversion			//
//--------------------------------------------------//

void ImageHandler::RGBtoLUV(Vec3b* rgbVal, float *luvVal)
{
	//delcare variables
	double	x, y, z, L0, u_prime, v_prime, constant;

	//convert RGB to XYZ...
	x		= XYZ[0][0]*(*rgbVal)[2] + XYZ[0][1]*(*rgbVal)[1] + XYZ[0][2]*(*rgbVal)[0];
	y		= XYZ[1][0]*(*rgbVal)[2] + XYZ[1][1]*(*rgbVal)[1] + XYZ[1][2]*(*rgbVal)[0];
	z		= XYZ[2][0]*(*rgbVal)[2] + XYZ[2][1]*(*rgbVal)[1] + XYZ[2][2]*(*rgbVal)[0];

	//convert XYZ to LUV...

	//compute L*
	L0		= y / (255.0 * Yn);
	if(L0 > Lt)
		luvVal[0]	= (float)(116.0 * (pow(L0, 1.0/3.0)) - 16.0);
	else
		luvVal[0]	= (float)(903.3 * L0);

	//compute u_prime and v_prime
	constant	= x + 15 * y + 3 * z;
	if(constant != 0){
		u_prime	= (4 * x) / constant;
		v_prime = (9 * y) / constant;
	}else
	{
		u_prime	= 4.0;
		v_prime	= 9.0/15.0;
	}

	//compute u* and v*
    luvVal[1] = (float) (13 * luvVal[0] * (u_prime - Un_prime));
    luvVal[2] = (float) (13 * luvVal[0] * (v_prime - Vn_prime));
}

//define inline rounding function...
inline int my_round(double in_x)
{
	if (in_x < 0)
		return (int)(in_x - 0.5);
	else
		return (int)(in_x + 0.5);
}

void ImageHandler::LUVtoRGB(float *luvVal, Vec3b* rgbVal)
{
	//declare variables...
	int		r, g, b;
	double	x, y, z, u_prime, v_prime;

	//perform conversion
	if(luvVal[0] < 0.1)
		r = g = b = 0;
	else
	{
		//convert luv to xyz...
		if(luvVal[0] < 8.0)
			y	= Yn * luvVal[0] / 903.3;
		else
		{
			y	= (luvVal[0] + 16.0) / 116.0;
			y  *= Yn * y * y;
		}

		u_prime	= luvVal[1] / (13 * luvVal[0]) + Un_prime;
		v_prime	= luvVal[2] / (13 * luvVal[0]) + Vn_prime;

		x		= 9 * u_prime * y / (4 * v_prime);
		z		= (12 - 3 * u_prime - 20 * v_prime) * y / (4 * v_prime);

		//convert xyz to rgb...
		//[r, g, b] = RGB*[x, y, z]*255.0
		r		= my_round((RGB[0][0]*x + RGB[0][1]*y + RGB[0][2]*z)*255.0);
		g		= my_round((RGB[1][0]*x + RGB[1][1]*y + RGB[1][2]*z)*255.0);
		b		= my_round((RGB[2][0]*x + RGB[2][1]*y + RGB[2][2]*z)*255.0);

		//check bounds...
		if(r < 0)	r = 0; if(r > 255)	r = 255;
		if(g < 0)	g = 0; if(g > 255)	g = 255;
		if(b < 0)	b = 0; if(b > 255)	b = 255;
	}

	//assign rgb values to rgb vector rgbVal
	(*rgbVal)[0]	= b;
	(*rgbVal)[1]	= g;
	(*rgbVal)[2]	= r;
}
